package com.xxx.encrypt.filter;

import com.xxx.encrypt.rsa.RsaKeys;
import com.xxx.encrypt.service.RsaService;
import com.xxx.encrypt.service.RsaServiceImpl;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.util.Strings;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.concurrent.atomic.AtomicReference;

@Component
@Slf4j
public class DecodeGatewayFilter implements GlobalFilter, Ordered {
    public static final String CACHE_REQUEST_BODY_OBJECT_KEY = "cachedRequestBodyObject";

    /**
     * 用于获取请求参数
     *
     * @param body
     * @return
     */
    private static String toRaw(Flux<DataBuffer> body) {
        AtomicReference<String> rawRef = new AtomicReference<>();
        body.subscribe(buffer -> {
            byte[] bytes = new byte[buffer.readableByteCount()];
            buffer.read(bytes);
            DataBufferUtils.release(buffer);
            rawRef.set(Strings.fromUTF8ByteArray(bytes));
        });
        return rawRef.get();
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest oldRequest = exchange.getRequest();
        URI uri = oldRequest.getURI();
        String method = oldRequest.getMethodValue();

        if ("POST".equals(method)){
            //获取request body
            Flux<DataBuffer> cachedBody = exchange.getAttribute(CACHE_REQUEST_BODY_OBJECT_KEY);
            if (cachedBody != null) {
                String raw = toRaw(cachedBody);
                log.info("request body is:{}", raw);
                //Todo 获取到RequestBody 做你想做的操作
                RsaService rsaService = new RsaServiceImpl();
                String decryptData = null;
                try {
                    decryptData = rsaService.decryptData(raw, RsaKeys.getServerPrvKeyPkcs8());
                } catch (Exception e) {
                    e.printStackTrace();
                }
                byte[] decryptBytes = decryptData.getBytes();

                // 根据解密后的参数重新构建请求
                DataBufferFactory dataBufferFactory = exchange.getResponse().bufferFactory();
                Flux<DataBuffer> bodyFlux = Flux.just(dataBufferFactory.wrap(decryptBytes));
                ServerHttpRequest newRequest = oldRequest.mutate().uri(uri).build();
                newRequest = new ServerHttpRequestDecorator(newRequest) {
                    @Override
                    public Flux<DataBuffer> getBody() {
                        return bodyFlux;
                    }
                };

                // 构建新的请求头
                HttpHeaders headers = new HttpHeaders();
                headers.putAll(exchange.getRequest().getHeaders());
                // 由于修改了传递参数，需要重新设置CONTENT_LENGTH，长度是字节长度，不是字符串长度
                int length = decryptBytes.length;
                headers.remove(HttpHeaders.CONTENT_LENGTH);
                headers.setContentLength(length);
                headers.set(HttpHeaders.CONTENT_TYPE, "application/json;charset=utf-8");
                newRequest = new ServerHttpRequestDecorator(newRequest) {
                    @Override
                    public HttpHeaders getHeaders() {
                        return headers;
                    }
                };
                return chain.filter(exchange.mutate().request(newRequest).build());
            }
        }

        return chain.filter(exchange);
    }

    @Override
    public int getOrder() {
        return -99;
    }
}
 